package com.mmg.cosplay.store;

import cn.dev33.satoken.stp.StpUtil;
import com.mmg.cosplay.entity.ChatMessages;
import com.mmg.cosplay.entity.po.RolePreset;
import com.mmg.cosplay.service.IRolePresetService;
import com.mmg.cosplay.service.RankingListService;
import com.mmg.cosplay.utils.RoleContext;
import com.mmg.cosplay.utils.TimestampUtils;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.ChatMessageDeserializer;
import dev.langchain4j.data.message.ChatMessageSerializer;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.data.mongodb.core.MongoTemplate;
import org.springframework.data.mongodb.core.query.Criteria;
import org.springframework.data.mongodb.core.query.Query;
import org.springframework.data.mongodb.core.query.Update;
import org.springframework.stereotype.Repository;

import java.util.LinkedList;
import java.util.List;

@Slf4j
@Repository
public class MongoChatMemoryStore implements ChatMemoryStore {

    // 辅助方法：提取第一条用户消息作为标题
    private static final String TEMP_TITLE_PREFIX = "[TEMP]";

    @Resource
    private MongoTemplate mongoTemplate;
    @Resource
    private IRolePresetService rolePresetService;

    @Resource
    private RankingListService rankingListService;

    @Override
    public List<ChatMessage> getMessages(Object memoryId) {
        Criteria criteria = Criteria.where("memoryId").is(memoryId.toString());
        Query query = new Query(criteria);
        ChatMessages chatMessages = mongoTemplate.findOne(query, ChatMessages.class);
        if (chatMessages == null) {
            return new LinkedList<>();
        }
        String contentJson = chatMessages.getContent();
        return ChatMessageDeserializer.messagesFromJson(contentJson);
    }

    @Override
    public void updateMessages(Object memoryId, List<ChatMessage> list) {
        Criteria criteria = Criteria.where("memoryId").is(memoryId.toString());
        Query query = new Query(criteria);
        ChatMessages existingMessage = mongoTemplate.findOne(query, ChatMessages.class);

        if (existingMessage == null) {
            // 如果不存在，创建新的文档
            ChatMessages newMessage = new ChatMessages();
            newMessage.setMemoryId(memoryId.toString());
            newMessage.setContent(ChatMessageSerializer.messagesToJson(list));

            // 设置用户id
            newMessage.setUserId(Long.parseLong(StpUtil.getLoginId().toString()));

            // 设置时间戳
            newMessage.setTimestamp(System.currentTimeMillis());

            Long roleId = RoleContext.getCurrentId();

            // 设置角色id
            newMessage.setRoleId(roleId);

            // 新会话时，角色使用次数+1（总计数）
            incrementRoleUsageCount(roleId);

            // 新会话时，排行榜计数+1（今日计数）
            incrementRankingUsageCount(roleId);

            // 设置时间
            newMessage.setTime(TimestampUtils.format(newMessage.getTimestamp()));

            // 提取第一个用户消息作为标题
            String title = extractFirstUserMessage(list);
            newMessage.setTitle(title != null ? title : TEMP_TITLE_PREFIX + "新对话"); // 使用特殊前缀标识临时标题

            // 设置默认不置顶
            newMessage.setIsTop("0");

            mongoTemplate.save(newMessage);

            // 清除上下文中的角色ID
            RoleContext.removeCurrentId();
        } else {
            // 如果存在，更新内容和时间
            Update update = new Update();
            update.set("content", ChatMessageSerializer.messagesToJson(list));
            // 如果当前标题是临时标题且现在有了用户消息，则更新标题
            if (existingMessage.getTitle() != null && existingMessage.getTitle().startsWith(TEMP_TITLE_PREFIX)) {
                String newTitle = extractFirstUserMessage(list);
                if (newTitle != null) {
                    update.set("title", newTitle);
                }
            }
            mongoTemplate.updateFirst(query, update, ChatMessages.class);
        }
    }

    /**
     * 增加角色使用次数（总计数）
     *
     * @param roleId 角色ID
     */
    private void incrementRoleUsageCount(Long roleId) {
        try {
            // 使用MyBatis-Plus的update方法来增加使用次数
            rolePresetService.lambdaUpdate()
                    .eq(RolePreset::getId, roleId)
                    .setSql("usage_count = IFNULL(usage_count, 0) + 1")
                    .update();
            log.info("角色ID: {} 总使用次数已增加", roleId);
        } catch (Exception e) {
            log.error("更新角色总使用次数失败，角色ID: {}", roleId, e);
        }
    }

    /**
     * 增加排行榜使用次数（今日计数）
     *
     * @param roleId 角色ID
     */
    private void incrementRankingUsageCount(Long roleId) {
        try {
            rankingListService.incrementRoleUsage(roleId);
            log.info("角色ID: {} 排行榜使用次数已增加", roleId);
        } catch (Exception e) {
            log.error("更新角色排行榜使用次数失败，角色ID: {}", roleId, e);
        }
    }

    @Override
    public void deleteMessages(Object memoryId) {
        Criteria criteria = Criteria.where("memoryId").is(memoryId.toString());
        Query query = new Query(criteria);
        mongoTemplate.remove(query, ChatMessages.class);
    }

    private String extractFirstUserMessage(List<ChatMessage> messages) {
        for (ChatMessage msg : messages) {
            if (msg instanceof UserMessage) {
                String text = ((UserMessage) msg).singleText();
                // 限制标题长度，避免过长
                return text.length() > 30 ? text.substring(0, 30) + "..." : text;
            }
        }
        // 第一次执行时没有用户消息，返回null，等待后续更新
        return null;
    }
}
